# -*- coding:utf-8 -*-#
# @Time:2023/7/7 21:21
# @Author:Adong
# @Software:PyCharm

import numpy as np
import os
import scipy.io.wavfile as _wavefile
import random
import numpy.random
import librosa
import matplotlib.pyplot as plt


def img(y,list):
    plt.plot(range(len(y)), y, label='y')
    for idx,i in enumerate(list):
        plt.plot(range(len(i)), i)
    # 简单的设置legend(设置位置)位置在右上角
    plt.legend(loc='upper right')
    plt.xlabel("time")  # x轴上的名字
    plt.ylabel("amplitude")  # y轴上的名字
    # plt.savefig(r'./result_img/' + model_name.split('/')[-1].split('.')[0] + '.png')
    plt.show()



def array2wav(nparray, framerate: int,wav_file: str):
    """
    将numpy数组转为单通道wav文件
    :param nparray: 输入的numpy向量
    :param wav_file: wav文件名
    :param framerate: 采样率(默认=sr)
    :return:
    """
    _wavefile.write(wav_file, framerate, nparray.astype(np.float32))



class wav_enhancement:
    def __init__(self,y,sr):
        self.y = y
        self.sr = sr

    def rolling(self,mode,savepath):
        """
        时域信号翻转
        :param mode:   y:纵向翻转;x:横向翻转
        :return:
        """
        y = self.y
        sr = self.sr
        if mode == 'y':
            y = -y
            array2wav(y,sr,savepath)
        elif mode == 'x':
            y = y[::-1]
            array2wav(y,sr,savepath)


    def mask(self,mode,savepath):
        """
        时域信号掩蔽
        :param mode:   random:随机掩蔽;successive:连续掩蔽
        :return:
        """
        y = self.y
        sr = self.sr
        if mode == 'random':
            mask = np.random.randint(0,2,len(y))
            y = y * mask
            array2wav(y,sr,savepath)
        elif mode == 'successive':
            start = np.random.randint(0,len(y))
            length = int(0.1 * len(y))
            mask = np.ones(len(y))
            mask[start:start+length] = 0
            y = y * mask
            array2wav(y,sr,savepath)

    def translate(self,savepath):
        '''
        时域信号平移
        :return:
        '''
        y = self.y
        sr = self.sr
        max = np.max(y)
        min = np.min(y)
        stride = random.uniform(min,max)
        y = y + stride
        array2wav(y, sr, savepath)

    def scale(self,savepath):
        '''
        时域信号缩放
        :return:
        '''
        y = self.y
        sr = self.sr
        scale = random.uniform(0, 5)
        y = y * scale
        array2wav(y, sr, savepath)

    def tfmask(self,mode,savepath):
        """
        时频掩蔽
        :param mode:
        time_random:时域随机掩蔽;time_successive:时域10%连续掩蔽;
        frequency_random:频域随机掩蔽;frequency_successive:频域10%连续掩蔽;
        :return:
        """
        y = self.y
        sr = self.sr
        n_fft = int(25 / 1000 * sr)
        hop_length = int(10 / 1000 * sr)
        stft = librosa.stft(y, n_fft=n_fft, hop_length=hop_length, window="hamm")  # stft得到时频谱

        # librosa.display.specshow(np.log(abs(stft)), sr=sr,n_fft=n_fft, hop_length=hop_length, x_axis='time', y_axis='log')
        if mode == 'time_random':
            mask = np.random.randint(0,2,len(stft[0]))
            stft = stft * mask
            istft = librosa.istft(stft, n_fft=n_fft, hop_length=hop_length, window="hamm")
            array2wav(istft,sr,savepath)
        elif mode == 'time_successive':
            length = int(0.1 * len(stft[0]))
            start = np.random.randint(0, len(stft[0])-length)
            mask = np.ones(len(stft[0]))
            mask[start:start + length] = 0
            stft = stft * mask
            istft = librosa.istft(stft, n_fft=n_fft, hop_length=hop_length, window="hamm")
            array2wav(istft, sr, savepath)
        elif mode == 'frequency_random':
            mask = np.random.randint(0, 2, len(stft))
            stft = stft.T * mask
            istft = librosa.istft(stft.T, n_fft=n_fft, hop_length=hop_length, window="hamm")
            array2wav(istft, sr, savepath)
        elif mode == 'frequency_successive':
            length = int(0.1 * len(stft))
            start = np.random.randint(0, len(stft)-length)
            mask = np.ones(len(stft))
            mask[start:start + length] = 0
            stft = stft.T * mask
            istft = librosa.istft(stft.T, n_fft=n_fft, hop_length=hop_length, window="hamm")
            array2wav(istft, sr, savepath)



if __name__ == '__main__':
    # filepath = r'data/wav_data_V3_train/normal/讲话声-TASCAM_0652_02.wav'
    # save = r'data/wav_data_V3_train/normal/'

    # n = len(os.listdir(save))  # 记录文件夹中已经存在的图片数量，避免重复生成覆盖
    # y,sr = librosa.load(filepath)
    # enhance = wav_enhancement(y,sr)

    # enhance.rolling('x',save + str(n) + '.wav')
    # enhance.rolling('y',save + str(n+1) + '.wav')

    # for i in range(1):
    #     n = len(os.listdir(save))  # 记录文件夹中已经存在的图片数量，避免重复生成覆盖
    #     enhance.mask('random', save + str(n) + '.wav')
    #     enhance.mask('successive', save + str(n + 1) + '.wav')
    #     enhance.translate(save + str(n + 2) + '.wav')
    #     enhance.scale(save + str(n + 3) + '.wav')
    #     enhance.tfmask('time_random',save + str(n + 4) + '.wav')
    #     enhance.tfmask('time_successive',save + str(n + 5) + '.wav')
    #     enhance.tfmask('frequency_random',save + str(n + 6) + '.wav')
    #     enhance.tfmask('frequency_successive',save + str(n + 7) + '.wav')
    filepath = r'C:\Users\83543\Desktop\新建文件夹\taishou-voiceprint-diagnosis\data\wav_data_V1\test\partial_discharge\局部放电_byq_jbfd_01.wav'
    y,sr = librosa.load(filepath)
    y_train = y[:int(0.6*len(y))]
    y_test = y[int(0.6*len(y)):int(0.8*len(y))]
    y_veri = y[int(0.8*len(y)):]
    ft = filepath.split('\\')[-2]
    _wavefile.write(f'{ft}_train.wav', sr, y_train.astype(np.float32))
    _wavefile.write(f'{ft}_test.wav', sr, y_test.astype(np.float32))
    _wavefile.write(f'{ft}_veri.wav', sr, y_veri.astype(np.float32))